using DataFrames
using QuadGK
using Distributions
using CSV
using LinearAlgebra
using JLD
using Interpolations
using SparseArrays
using DelimitedFiles

clearconsole()

#-------------------------------------------------------------
# IRF Configuration
#-------------------------------------------------------------

H = 60 # maximum horizon
n_every = 50 # use every n_every'th draw from the posterior
sh_size = 3   # shock size, in multiples of standard deviations
sh_id = 1 # sh_id = 1: TFP shock

#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
include(readDir *"vech.jl");
include(readDir *"logSpline_Procedures.jl");
include(readDir *"VAR_Procedures.jl");
include(readDir *"Loaddata.jl");
include(readDir *"VAR_MoreLags_Procedures.jl")

#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nDataSel  = "1"
nfVARSpec = "3"
nModSpec  = "3"
nMCMCSpec = "1"
modName   = "SS"  # VAR or SS

specDir   = "$(pwd())/SpecFiles/"
include(specDir * "/fVARspec" * nfVARSpec * ".jl")
include(specDir * "/" * modName * "spec" * nModSpec * ".jl")
include(specDir * "/" * modName * "MCMCspec" * nMCMCSpec * ".jl")

#-------------------------------------------------------------
# load aggregate data
#-------------------------------------------------------------
dataDir  = "$(pwd())/data/Para" * nDataSel *"/"
agg_data = loadaggdata(1,Tend)
n_agg    = size(agg_data)[2]


#-------------------------------------------------------------
# Load coefficients from density estimation and knots
#-------------------------------------------------------------
sNameLoadDir = "fVAR" * nfVARSpec
loaddir      = "$(pwd())/results/Para" *nDataSel * "/" * sNameLoadDir *"/";
sNameFile    = "K" * string(K) * "_fVAR" * nfVARSpec

PhatDensCoef_factor, MDD_GoF, VinvLam_all, PhatDensCoef_lambda, PhatDensCoef_mean  = loaddensdata(1,Tend,K,nfVARSpec)
n_cross = size(PhatDensCoef_factor)[2]

knots_all = readdlm(loaddir * "fVAR" *nfVARSpec * "_knots_all.csv", ',', Float64, '\n'; skipstart = 1);
ii=getindex.(findall(K_vec.-K.==0),[1 2])[1] # find index ii where K==K_vec
knots = knots_all[quant_sel[ii,:].==1]

# knots=loadknots(quant_vec, quant_sel)

#-------------------------------------------------------------
# Adjust sample
#-------------------------------------------------------------

# Arrange data for VAR analysis
# We start by dropping the initial Tdrop-nlags observations
agg_data_adj = agg_data[Tdrop-nlags+1:end,:]
PhatDensCoef_factor_adj = PhatDensCoef_factor[Tdrop-nlags+1:end,:]
# We now have T*=(T-T0)+nlags observations and call the first observation t=1
# Observations t=1,...,nlags are used to initialize lags
# and observations t=nlags+1,...,T* are the estimation sample
# Note that by holding Tdrop fixed and changing nlags<=Tdrop we can change
# the nunber of lags holding the estimation sample fixed.
# Estimation sample has T*-nlags=T-T0 observations

# Define some constants
n       = n_agg+n_cross
p       = nlags

# Combine aggregate and cross sectional observations and create YY XX for VAR estimation
data_adj = [ agg_data_adj PhatDensCoef_factor_adj ]
YY, XX, PHIols, Sols, SIGMAols = data_to_YY(data_adj, nlags, nlags)
# FS: having nlags twice as input arguments is not necessary. We can simplify
# but would have to adjust other programs as well.



#-------------------------------------------------------------
# Load posterior draws
#-------------------------------------------------------------

sName    = "fVAR" * nfVARSpec * "_" * modName * nModSpec * "_" * "MCMC" * nMCMCSpec;
loadDir  = "$(pwd())/results/Para" *nDataSel * "/" * sName *"/";

Bpmean     = load(loadDir * sName * "_PostMeans.jld", "Bpmean")
Sigpmean   = load(loadDir * sName * "_PostMeans.jld", "Sigpmean")
Sigtrpmean = load(loadDir * sName * "_PostMeans.jld", "Sigtrpmean")

Bpdraw     = load(loadDir * sName * "_PostDraws.jld", "Bpdraw")
Sigpdraw   = load(loadDir * sName * "_PostDraws.jld", "Sigpdraw")
Sigtrpdraw = load(loadDir * sName * "_PostDraws.jld", "Sigtrpdraw")

#-------------------------------------------------------------
# Generate IRFs at posterior mean
#-------------------------------------------------------------

println("")
println("Generating IRFs at posterior mean... ")
println("")

YY_IRF          = zeros(H+1,n)
L               = Sigtrpmean
Btilde          = reshape(Bpmean, n*p,n)
response        = IRFredu(Btilde,L,H,sh_size,sh_id)
YY_IRF[2:end,:] = response'

PhatDens_IRF    = zeros(H+1,length(xgrid));
# compute density IRFs
for hh = 1:H+1
    PhatDensCoef_hh    = coefRecover(YY_IRF[hh,n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
    PhatDensNorm_hh    = lnpdfNormalize(PhatDensCoef_hh,knots,minimum(xgrid), maximum(xgrid))
    PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
end

savedir = "$(pwd())/results/Para" *nDataSel * "/" * sName *"/";
try mkdir(savedir) catch; end
CSV.write(savedir * sName * "_IRF_YY_AggSh"*string(sh_id)*"_pmean.csv", DataFrame(YY_IRF,:auto))
CSV.write(savedir * sName * "_IRF_PhatDens_AggSh"*string(sh_id)*"_pmean.csv", DataFrame(PhatDens_IRF,:auto))


#-------------------------------------------------------------
# Generate IRFs for a subset of posterior draws
#-------------------------------------------------------------

println("")
println("Generating IRFs for posterior draws... ")
println("")

nsim = size(Bpdraw)[1]
n_subseq                 = floor(Int,nsim/n_every)
YY_IRF_uncertainty       = zeros(H+1,n,n_subseq)

# hh=1 is steady state
# hh=2 is period of impact

for pp = 1:n_subseq

    time_init_loop = time_ns();
    println(" Draw number:  $pp")
    println(" Remaining draws:  $(n_subseq-pp)")
    println(" Posterior draw number: $(pp*n_every)")

    L      = Sigtrpdraw[pp*n_every,:,:]
    Btilde = Bpdraw[pp*n_every,:,:]
    response = IRFredu(Btilde,L,H,sh_size,sh_id)
    YY_IRF_uncertainty[2:end,:,pp] = response'

    PhatDens_IRF   = zeros(H+1,length(xgrid));
    #Gini_IRF       = zeros(H+1,1)

    # compute density IRFs
    for hh = 1:H+1
        PhatDensCoef_hh    = coefRecover(YY_IRF_uncertainty[hh,n_agg+1:end,pp]', PhatDensCoef_lambda, PhatDensCoef_mean)
        PhatDensNorm_hh    = lnpdfNormalize(PhatDensCoef_hh,knots, minimum(xgrid), maximum(xgrid))
        PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
    end

    CSV.write(savedir * sName * "_IRF_YY_AggSh"*string(sh_id)*"_" * string(pp) * ".csv", DataFrame(YY_IRF_uncertainty[:,:,pp]))
    CSV.write(savedir * sName * "_IRF_PhatDens_AggSh" * string(sh_id) * "_" * string(pp) * ".csv", DataFrame(PhatDens_IRF))

    time_loop=signed(time_ns()-time_init_loop)/1000000000
    println(" Elapsed time = $(time_loop)")
    println("")

end


#-------------------------------------------------------------
# Impulse responses based on KS-VAR
#-------------------------------------------------------------

#println("")
#println("Generating IRFs for KS-VAR... ")
#println("")


#measureCoef_2_1 = simul_data[1:Tend,8];
#measureCoef_2_2 = simul_data[1:Tend,9];
#measureCoef_2_3 = simul_data[1:Tend,10];

#moment_2_1 = simul_data[1:Tend,16];
#moment_2_2 = simul_data[1:Tend,17];
#moment_2_3 = simul_data[1:Tend,18];

#PDensCoef = [measureCoef_2_1 measureCoef_2_2 measureCoef_2_3 moment_2_1 moment_2_2 moment_2_3]
#(PDensCoef_factor, PDensCoef_lambda, PDensCoef_mean) = coefCompress(PDensCoef);

#aggEmployment = 0.93 # from Winberry's calibration

#KSVAR_data    = [agg_data PDensCoef_factor]

#~, KSVAR_XX, KSVAR_PHI, ~, KSVAR_SIGMA = data_to_YY(KSVAR_data, Tdrop, nlags)

#KSVAR_SIGMAtr = cholesky(Symmetric(KSVAR_SIGMA)).L

#nx          = size(KSVAR_XX,2)
#YY_IRF      = zeros(H+1,nx)
#PDens_IRF   = zeros(H+1,length(xgrid))

# hh=1 is steady state
# hh=2 is period of impact

# impact
#YY_IRF[1+1,:]  = sh_size*KSVAR_SIGMAtr[:,1]'

# future periods
#for hh = 1+2:H+1
#    YY_IRF[hh,:]  = YY_IRF[hh-1,:]'*KSVAR_PHI
#end

# first column is Zt
#Agg_IRF = YY_IRF[:,1:n_agg]

# columns characterize the density
#for hh = 1:H+1
#    PDensCoef_hh    = coefRecover(YY_IRF[hh,n_agg+1:end]', PDensCoef_lambda, PDensCoef_mean)
#    PDens_IRF[hh,:] = pdfWinberry(PDensCoef_hh[1:3]',PDensCoef_hh[4],PDensCoef_hh[5],PDensCoef_hh[6],xgrid, aggEmployment)
#end

#CSV.write(savedir * sName * "_PDens_IRF.csv", DataFrame(PDens_IRF))
#CSV.write(savedir * sName * "_PAgg_IRF.csv", DataFrame(Agg_IRF))
#CSV.write(savedir * sName * "_PYY_IRF.csv", DataFrame(YY_IRF))
